package com.huawei.codearts.tools.ratelimit.aspect;


import com.google.common.util.concurrent.RateLimiter;
import com.huawei.codearts.tools.ratelimit.annotation.RateLimit;
import com.huawei.codearts.tools.ratelimit.annotation.RateLimitCustom;
import com.huawei.codearts.tools.ratelimit.service.TokenBucketLimiter;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;


@Aspect
@Component
public class RateLimitAspect {

    private final Map<String, RateLimiter> limiters = new ConcurrentHashMap<>();
    private final ConcurrentMap<String, TokenBucketLimiter> limitersCustom = new ConcurrentHashMap<>();

    @Before("@annotation(rateLimit)")
    public void rateLimit(JoinPoint joinPoint, RateLimit rateLimit) throws Throwable {
//        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
//        Method method = signature.getMethod();
//        String methodName = method.getName();
        String methodName = joinPoint.getSignature().toShortString();
        System.out.println("rateLimit key..."+methodName);
        RateLimiter rateLimiter = limiters.computeIfAbsent(methodName, k -> RateLimiter.create(rateLimit.value()));
        if (!rateLimiter.tryAcquire()) {
            throw new RuntimeException("接口访问太过频繁，请稍候再试");
        }
    }

    @Before("@annotation(rateLimitCustom)")
    public void rateLimitCustom(JoinPoint joinPoint, RateLimitCustom rateLimit) throws Throwable {
        String methodName = joinPoint.getSignature().toShortString();
        System.out.println("rateLimitCustom key..."+methodName);
        TokenBucketLimiter rateLimiter = limitersCustom.computeIfAbsent(methodName, k ->
                new TokenBucketLimiter(rateLimit.maxTokens(), rateLimit.refillRate())
        );
        if (!rateLimiter.tryAcquire()) {
            throw new RuntimeException("接口访问太过频繁，请稍候再试!");
        }
    }

}


